package com.gome.gop.mybatis.rw.datasource;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import javax.sql.DataSource;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import com.alibaba.druid.pool.DruidDataSource;

/**
 * 根据MyDynamicDataSourceHolder中的线程局部变量，选择数据源（读/写） Created by mahengyang on
 * 2017/7/21.
 */
public class RoutingDataSource extends AbstractRoutingDataSource {
    private DruidDataSource writeDataSource; // 写数据源
    
    private List<DruidDataSource> readDataSources; // 读数据源
    
    private List<DruidDataSource> multiDataSources;
    
    private AtomicInteger count = new AtomicInteger(0);
    
    public RoutingDataSource(DruidDataSource writeDataSource, List<DruidDataSource> readDataSources, List<DruidDataSource> multiDataSources) {
        this.writeDataSource = writeDataSource;
        this.readDataSources = readDataSources;
        this.multiDataSources = multiDataSources;
        
    }
    
    /**
     * 在初始化之前被调用，设置默认数据源，以及数据源资源（这里的写法是参考源码中的）
     */
    @Override
    public void afterPropertiesSet() {
        // 如果写数据源不存在，则抛出非法异常
        if (this.writeDataSource == null) {
            throw new IllegalArgumentException("Property 'writeDataSource' is required");
        }
        // 设置默认目标数据源为主库
        setDefaultTargetDataSource(writeDataSource);
        // 设置所有数据源资源，有从库添加，没有就添加
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DynamicDataSourceGlobal.WRITE.getType(), writeDataSource);
        // 从库
        if (readDataSources != null && !readDataSources.isEmpty()) {
            for (int i = 0; i < readDataSources.size(); i++) {
                targetDataSources.put(this.readDataSourceName(i), readDataSources.get(i));
            }
        }
        // 其他手动切换的数据库
        if (multiDataSources != null && !multiDataSources.isEmpty()) {
            for (DruidDataSource dataSource : multiDataSources) {
                // multiDataSourceMap.put(dataSource.getName(), dataSource);
                targetDataSources.put(dataSource.getName(), dataSource);
            }
        }
        setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
    
    private String readDataSourceName(int index) {
        return DynamicDataSourceGlobal.READ.getType() + "_" + index;
    }
    
    /**
     * 这是AbstractRoutingDataSource类中的一个抽象方法，而它的返回值是你所要用的数据源dataSource的key值
     */
    @Override
    protected Object determineCurrentLookupKey() {
        // 根据当前线程所使用的数据源进行切换
        DynamicDataSourceGlobal dynamicDataSourceGlobal = DynamicDataSourceHolder.getDataSource();
        // 如果没有被赋值，那么默认使用主库
        if (dynamicDataSourceGlobal == null || dynamicDataSourceGlobal == DynamicDataSourceGlobal.WRITE) {
            return DynamicDataSourceGlobal.WRITE.getType();
        }
        // 判断是否为从库，如果为从库则，随机读取其中的一个数据源
        if (dynamicDataSourceGlobal == DynamicDataSourceGlobal.READ) {
            int index = Math.abs(count.incrementAndGet()) % readDataSources.size();
            return this.readDataSourceName(index);
        }
        // 其他情况使用从库
        return dynamicDataSourceGlobal.getName();
    }
    
    public void setWriteDataSource(DruidDataSource writeDataSource) {
        this.writeDataSource = writeDataSource;
    }
    
    public Object getWriteDataSource() {
        return writeDataSource;
    }
}